Asymptotic Estimation of Weight RMS in AdamW (Part 2)

In the blog post "Asymptotic Estimation of Weight RMS in AdamW (Part 1)", we derived the asymptotic expression for the RMS of model weights trained with AdamW. However, we assumed fixed Weight Decay and learning rate throughout training, which does not fully align with practical training. Therefore, in this article, we extend the previous conclusions to a dynamic version.

The dynamic version allows both Weight Decay and learning rate to vary with increasing training steps, such as classic Cosine Decay, WSD (Warmup Stable Decay), etc., making the conclusions more general.

Step One#

Our starting point remains the definition of AdamW:

Since \(\eta_t\lambda_t\ll 1\), we can write:

Let \(\kappa_t = \sum_{i=1}^t \eta_i\lambda_i\), then direct expansion gives:

Then set \(z_t = \sum_{i=1}^t e^{\kappa_i}\eta_i\), and by mean-field approximation:

Thus:

Step Two#

Following the previous approach, to estimate \(\Vert \bar{\boldsymbol{u}}_t\Vert_{RMS}^2\), we need to assume \(\boldsymbol{g}_j\) are i.i.d. from \(\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\sigma}^2)\), then compute:

Finally, averaging over components of \(\mathbb{E}[\bar{\boldsymbol{u}}_t^2]\) yields an approximation for \(\Vert\bar{\boldsymbol{u}}_t\Vert_{RMS}^2\).

Expanding \(\boldsymbol{m}_t,\boldsymbol{v}_t\) gives:

We also have the identity:

Using these results, we can write:

Step Three#

First compute the denominator. When \(t\) is sufficiently large (\(\beta_1^t, \beta_2^t\) sufficiently small), \(\sum_{j=1}^t \bar{\beta}_1(j,t)\) and \(\sum_{j=1}^t \bar{\beta}_2(j,t)\) are sufficiently close to 1, so:

Similarly, \(\mathbb{E}[\bar{\boldsymbol{m}}_t] = \boldsymbol{\mu}\), and \(\mathbb{E}[\bar{\boldsymbol{m}}_t^2] = \mathbb{E}[\bar{\boldsymbol{m}}_t]^2 + \mathbb{V}ar[\bar{\boldsymbol{m}}_t]\). Using additivity of variance:

Thus:

And:

Finally:

If readers are viewing this article directly, some steps may appear to jump; in such cases, it is advisable to revisit "Asymptotic Estimation of Weight RMS in AdamW (Part 1)" to familiarize with the reasoning behind each approximation.

Example One#

First consider \(\boldsymbol{\mu}=\boldsymbol{0}\), substituting the expression for \(\bar{\beta}_1(j,t)\) into the above equation yields:

Now consider the simple case \(\lambda_t=0\), i.e., no Weight Decay, then:

If \(\beta_1\to 0\), we immediately get \(\Vert\boldsymbol{\theta}_t\Vert_{RMS}^2 \approx \Vert\boldsymbol{\theta}_0\Vert_{RMS}^2 + \sum_{j=1}^t\eta_j^2\). This indicates that without Weight Decay and as training steps \(t\to\infty\), to prevent Weight RMS from exploding, the sum of squares of the learning rate sequence must converge.

In fact, even for \(0 < \beta_1 < 1\), this condition remains necessary and sufficient:

The proof is not difficult. Transform the left side:

This shows that if the left side converges, then for all \(i_1, i_2\), the sum \(\sum_{j=1}^{\infty}\eta_{i_1+j}\eta_{i_2+j}\) converges, which naturally implies \(\sum_{j=1}^{\infty}\eta_j^2\) converges, proving necessity. For sufficiency, starting from the above and using Cauchy-Schwarz inequality:

Thus convergence of \(\sum_{j=1}^{\infty}\eta_j^2\) implies convergence of the left side, proving sufficiency.

Example Two#

Next, we consider the case where Weight Decay is constant but learning rate varies, i.e., \(\kappa_t = \lambda\sum_{i=1}^t \eta_i\). If we want infinite training to obtain a solution as close as possible to theoretical optimum, the learning rate should satisfy \(\sum_{i=1}^{\infty} \eta_i \to \infty\).

For the general case, computing Equation (15) is challenging, but we can consider further approximations based on practical scenarios. In actual training, typically \(\lambda_t \eta_t \ll 1\), so \(e^{\kappa_i}\) grows much slower than \(\beta_1^i\) decays. Thus, we can approximate:

Substituting this approximation back into Equation (15):

Now we can only compute specifically for given \(\eta_j\). For example, when \(\lambda_j,\eta_j\) are constants, we have \(\kappa_t = \lambda\eta t\), and:

This matches the result from the previous article.

Differential Equation#

Considering integrals are usually easier to compute than summations, we can approximate sums with integrals:

where \(\kappa_t = \int_0^t \lambda_s\eta_s ds\). Let \(\rho_t = \Vert\boldsymbol{\theta}_t\Vert_{RMS}^2\), multiply both sides by \(e^{2\kappa_t}\) and differentiate:

This is the differential equation governing the RMS squared. If \(\rho_t\) converges to a constant as \(t\to\infty\), the left side equals 0, giving:

This tells us that for decay-type learning rate schedules, the final learning rate should not be set to 0; otherwise, prolonged training risks model weight collapse.

Mean Field#

The scenario where \(t\to\infty\) typically applies to multi-epoch supervised training. In pre-training scenarios, training is usually single-epoch, and \(\kappa_t\) is often \(\mathcal{\Theta}(1)\).

Under the assumption \(\kappa_t=\mathcal{\Theta}(1)\), we can consider mean-field approximation. Starting from the integral form (23), by definition, \(\kappa_s\) is a monotonically increasing function from 0 to \(\kappa_t\) over \([0,t]\). Thus:

That is, the target integral itself is bounded between \(e^{- 2\kappa_t} \nu_t\) and \(\nu_t\), where \(\nu_t = \int_0^t \eta_s^2 ds\). When \(\kappa_t=\mathcal{\Theta}(1)\), \(e^{-2\kappa_t}\) is not much smaller than 1, meaning \(\nu_t\) itself may be a good approximation.

We can be more precise by estimating a reasonable multiplier:

Substituting into Equation (23):

Example Three#

Returning to our main concern—the common setting of "fixed Weight Decay, variable learning rate"—we compute \(\kappa_t,\nu_t\) for several specific examples.

First, linear learning rate:

Integrating:

Next, Cosine Decay:

Integrating:

Finally, WSD (Warmup Stable Decay):

We have:

Simulation#

We can also validate the above approximations through numerical simulation:

import numpy as np

N, T = 10000, 10000
beta1, beta2 = 0.9, 0.95
m, v = 0, 0
w = np.random.randn(N) * (init_std := 0.1)
lr_max, lr_min, wd = 0.001, 0.0001, 0.1
lr = lr_min + (lr_max - lr_min) * (1 + np.cos(np.arange(T) / T * np.pi)) / 2
for i in range(T):
    g = np.random.randn(N)
    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g**2
    w = w - lr[i] * (m / v**0.5 + wd * w)

# Direct computation ≈ 0.0744
weight_rms = (w**2).mean()**0.5

# Series approximation ≈ 0.0742
kappa = wd * lr.cumsum()
approx1 = ((np.exp(kappa * 2) * lr**2).sum() + init_std**2)**0.5 * np.exp(-kappa[-1])

# Mean-field approximation ≈ 0.0760
kappa = wd * (lr_max + lr_min) / 2 * T
nu = (3 * lr_max**2 + 2 * lr_max * lr_min + 3 * lr_min**2) / 8 * T
approx2 = ((np.exp(kappa * 2) - 1) * nu / kappa / 2 + init_std**2)**0.5 * np.exp(-kappa)

print(weight_rms)
print(approx1)
print(approx2)

Summary#

This article extends the results from Part 1 to a dynamic version, allowing us to estimate Weight RMS for AdamW with time-varying learning rates and Weight Decay.